'''
Author: zhanwei xu
Date: 2023-12-06 15:17:55
LastEditors: zhanwei xu
LastEditTime: 2023-12-07 10:08:25
Description: 

Copyright (c) 2023 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''

from .utils import *
from fastapi.responses import StreamingResponse, JSONResponse
from datetime import datetime
import threading
import time
import os
import random
import base64

current_time = datetime.now()
print(current_time)

def save_count(chatbot):
    while True:
        current_time = datetime.now()
        if current_time.hour == 1 and current_time.minute == 0: #检查时间是否凌晨1点
            with open('count.txt', 'a') as f: # 打开文件以追加的方式保存count
                f.write(str(current_time) + ': ' + str(chatbot.count) + '\n')  # 保存count
            chatbot.count = 0  # 重置count
            time.sleep(60)  # 休眠60秒，以避免重复保存
        else:
            time.sleep(10) # 休眠10秒，然后再次检查时间


class Visionbot():  
    def __init__(self) -> None:
        self.model_name = "gpt-3.5-turbo"
        self.ip_traffic = {}
        self.ip_traffic_4 = {}
        self.api_key_index = 0
        self.api_keys = get_api_keys('api_keys.txt')
        self.api_key_index = random.randint(0,400)
        self.count = 0
        self.count_thread = threading.Thread(target=save_count, args=(self,))
        self.count_thread.daemon = True # 设置为守护线程，这样在主程序退出时，线程也会自动退出
        self.count_thread.start() # 启动线程
    

    async def visionchat(self, request):
        """
        This is an async function that processes a chatbot request, checks if the key is valid and not
        empty, decrements the key count, generates a response based on the prompt, and returns a streaming
        response.
        
        :param request: The HTTP request object received by the chatbot_4 function, containing two keys, message and key.
        :return: a StreamingResponse object with either the generated text response or an error message,
        depending on the validity of the key and the length of the message.
        """
        flag = 0
        # client_ip = request.client.host
        client_ip = request.headers.get("X-Real-IP")
        json_data = await request.json()
        print(client_ip)
        if check_traffic_4(client_ip):
            try:
                key = json_data["key"]
            except:
                return StreamingResponse(generate_error_key(), media_type='application/json')
            key_list = os.listdir("keys")
            # check if the key is valid
            if (key not in key_list) or (os.stat("keys/" + key).st_size == 0):
                key_list = os.listdir('key_token')
                if (key not in key_list) or (os.stat("key_token/" + key).st_size == 0):
                    print('adafd')
                    return StreamingResponse(generate_error_key(), media_type='application/json')
                else:
                    flag = 1
            else:
                flag = 2
        
        
        # Extract message from request
        message = json_data["message"]
        


        response = await async_request_4(message)
        if response.status_code!=200:                 
            print(response.status_code)
            print(response.text)
            return StreamingResponse(generate_error(), media_type="text/plain")
        self.count += 1
        if flag == 0:
            return StreamingResponse(generate_text(response), media_type="text/event-stream")
        if flag == 1: # key_token
            token_num = num_tokens_from_message(message,model="gpt-4")
            consume_token(key,token_num)
            consume_token(key,1000)
            return StreamingResponse(generate_text(response,user_key=key), media_type="text/event-stream" )
        else: # key
            
            key_path = "keys/" + key
            with open(key_path, 'r') as f:
                key_conunt = f.read()
                print(key_conunt)
                key_conunt = int(key_conunt)
                key_count = key_conunt -1
            with open(key_path, 'w') as f:
                f.write(str(key_count))
            if key_count <= 0:
                os.remove(key_path)
        return StreamingResponse(generate_text(response), media_type="text/event-stream" )
